package edu.cornell.cs4740.postag;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Set;

import edu.cornell.cs4740.postag.parsing.Sentence;
import edu.cornell.cs4740.postag.parsing.WordContainer;


public class Baseline {
  // word, pos
  HashMap<String, String> mostFrequentPos = new HashMap<String, String>();
  List<Sentence> training;
  
  public Baseline(List<Sentence> training){
    this.training = training;
  }
  
  public List<Sentence> getPredictions(List<Sentence> test) {
    if (mostFrequentPos.size() == 0) {
      findMostFrequentPos();
    }
    
    List<Sentence> predictions = new ArrayList<Sentence>();
    for (int i = 0; i < test.size(); i++) {
      Sentence curSentence = test.get(i);
      List<WordContainer> curWords = curSentence.getAllWords();
      List<WordContainer> predictedWords = new ArrayList<WordContainer>();
      
      for (int j = 0; j < curWords.size(); j++) {
        WordContainer curWord = curWords.get(j);
        String pos = mostFrequentPos.get(curWord.getWord());
        if (pos == null) {
          pos = "";
        }
        WordContainer predictedWord = new WordContainer(curWord.getWord(), pos);
        predictedWords.add(predictedWord);
      }
      
      Sentence predictedSentence = new Sentence(predictedWords);
      predictions.add(predictedSentence);
    }
    
    return predictions;
  }
  
  private void findMostFrequentPos() {
    HashMap<String, HashMap<String, Integer>> allWordsAndPos = new HashMap<String, HashMap<String, Integer>>();
    for (int i = 0; i < training.size(); i++) {
      Sentence curSentence = training.get(i);
      
      for (int j = 0; j < curSentence.getLength(); j++) {
        WordContainer curWord = curSentence.getWord(j);
        
        HashMap<String, Integer> curPosCount = allWordsAndPos.get(curWord.getWord());
        if (curPosCount == null) {
          curPosCount = new HashMap<String, Integer>();
        }
        
        Integer curCount = curPosCount.get(curWord.getPos());
        if (curCount == null) {
          curCount = 1;
        } else {
          curCount += 1;
        }
        
        curPosCount.put(curWord.getPos(), curCount);
        allWordsAndPos.put(curWord.getWord(), curPosCount);
      }
    }
    
    Set<String> keyset = allWordsAndPos.keySet();
    for (String word : keyset) {
      HashMap<String, Integer> allPos = allWordsAndPos.get(word);
      String maxPos = "";
      int maxCount = 0;
      Set<String> allPosKeys = allPos.keySet();
      
      for (String posKey : allPosKeys) {
        int count = allPos.get(posKey);
        if (count > maxCount) {
          maxPos = posKey;
          maxCount = count;
        }
      }
      
      mostFrequentPos.put(word, maxPos);
    }
  }

}
